package com.xmzs.sd.service.impl;


import cn.hutool.core.bean.BeanUtil;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson2.JSONArray;
import com.xmzs.common.core.domain.model.LoginUser;
import com.xmzs.common.core.exception.base.BaseException;
import com.xmzs.common.oss.core.CosClient;
import com.xmzs.common.satoken.utils.LoginHelper;
import com.xmzs.sd.domain.SdModel;
import com.xmzs.sd.domain.bo.SysFileInfoBo;
import com.xmzs.sd.domain.enums.FileType;
import com.xmzs.sd.domain.request.Txt2ImgRequest;
import com.xmzs.sd.domain.response.Txt2ImgResponse;
import com.xmzs.sd.domain.vo.SdModelParamVo;
import com.xmzs.sd.service.ISdModelParamService;
import com.xmzs.sd.service.ISysFileInfoService;
import com.xmzs.sd.service.StableDiffusionService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.TimeUnit;


/**
 * 描述： StableDiffusion接口
 *
 * @author wangle
 * 2023-07-15
 */
@Service
@RequiredArgsConstructor
@Slf4j
public class StableDiffusionServiceImpl implements StableDiffusionService {

    private final ISysFileInfoService sysFileInfoService;

    private final CosClient cosClient;

    private final ISdModelParamService sdModelParamService;

    @Value("${sd.apiHost}")
    private String apiHost;

    private MediaType JSON = MediaType.get("application/json");

    /**
     * 查询模型列表
     *
     * @param
     * @return List<String>
     * @Date 2023/7/15
     **/
    @Override
    public List<SdModel> modelList() {
        Request request = new Request.Builder().url(this.apiHost + "/sdapi/v1/sd-models").build();
        String res = "";
        try {
            Response response = createClient().newCall(request).execute();
            if (response.isSuccessful()) {
                res = response.body().string();
            }
        } catch (IOException e) {
            throw new BaseException("未查询到模型信息！");
        }
        return JSONArray.parseArray(res, SdModel.class);
    }

    /**
     * 文生图
     *
     * @param prompt
     * @param modelName
     * @return String
     * @Date 2023/7/13
     **/
    @Override
    public List<String> txt2Img(String prompt, String modelName) {
        LoginUser loginUser = LoginHelper.getLoginUser();
        if (loginUser == null) {
            throw new BaseException("用户未登录!");
        }
        // 切换模型
        setOption(modelName);
        // 根据模型名称查询默认设置
        SdModelParamVo sdModelParamVo = sdModelParamService.queryByModelName(modelName);
        Txt2ImgRequest txt2ImgRequest = new Txt2ImgRequest();
        BeanUtil.copyProperties(sdModelParamVo, txt2ImgRequest);
        String json = JSONObject.toJSONString(txt2ImgRequest);
        // 构建请求体
        RequestBody body = RequestBody.create(JSON, json);
        Request request = new Request.Builder().post(body)
            .url(this.apiHost + "/sdapi/v1/txt2img").build();
        List<String> urlList = new ArrayList<>(10);
        try {
            Response response = createClient().newCall(request).execute();
            Txt2ImgResponse txt2ImgResponse = JSONObject.parseObject(response.body().string(), Txt2ImgResponse.class);
            List<String> imgBase64 = txt2ImgResponse.getImages();
            for (int i = 0; i < imgBase64.size(); i++) {
                String url = cosClient.uploadBase64(imgBase64.get(i));
                urlList.add(url);
                // 保存文件信息
                SysFileInfoBo bo = new SysFileInfoBo();
                bo.setUrl(url);
                bo.setFileType(FileType.FILE_SD.getFileType());
                bo.setUserId(loginUser.getUserId());
                // 保存图片信息
                sysFileInfoService.insertByBo(bo);
            }
        } catch (Exception e) {
            throw new BaseException("图片生成失败！");
        }
        return urlList;
    }

    /**
     * 设置参数
     *
     * @param modelName
     * @return void
     * @Date 2023/7/26
     **/
    @Override
    public void setOption(String modelName) {
        HashMap<Object, Object> map = new HashMap<>(10);
        map.put("sd_model_checkpoint", modelName);
        String json = JSONObject.toJSONString(map);
        // 设置模型
        RequestBody body = RequestBody.create(JSON, json);
        Request request = new Request.Builder().post(body).url(this.apiHost + "/sdapi/v1/options").build();
        String res;
        try {
            Response response = createClient().newCall(request).execute();
            res = response.body().string();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        log.info("{}参数设置成功{}",modelName,res);
    }

    /**
     * 创建一个默认的okHttp客户端
     *
     * @param
     * @return OkHttpClient
     * @Date 2023/7/26
     **/
    public OkHttpClient createClient() {
        OkHttpClient okHttpClient = new OkHttpClient
            .Builder()
            .connectTimeout(120, TimeUnit.SECONDS)
            .writeTimeout(600, TimeUnit.SECONDS)
            .readTimeout(600, TimeUnit.SECONDS)
            .build();
        return okHttpClient;
    }
}
